package com.chenbei.common.util;

import com.chenbei.common.domain.vo.TreeNode;
import io.swagger.v3.oas.annotations.media.Schema;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

@Schema(description = "树节点工具类")
public class TreeNodeUtil {

    /**
     * 组装树结构
     * @param all 数据源
     * @param rootNodeId 根节点ID
     * @return List<TreeNode<T>>
     * @param <T> 泛型
     */
    public static <T> List<TreeNode<T>> buildTree(List<TreeNode<T>> all, long rootNodeId) {
        if (all != null && !all.isEmpty()) {
            return all.stream()
                    .filter(tTreeNode -> tTreeNode.getPid() == rootNodeId)
                    .peek(tTreeNode -> tTreeNode.setChildren(getChildren(tTreeNode, all)))
                    .collect(Collectors.toList());
        }
        return null;
    }

    /**
     * 获取指定父节点的所有子节点集合
     * @param treeNodes 所有节点
     * @param pid 父ID
     * @return List
     * @param <T> 泛型
     */
    public static <T> List<TreeNode<T>> allChildrenNodes(List<TreeNode<T>> treeNodes, long pid) {
        if (treeNodes != null && !treeNodes.isEmpty()) {
            List<TreeNode<T>> childrenNodes = new ArrayList<>();
            for (TreeNode<T> treeNode : treeNodes) {
                if (treeNode.getPid() == pid) {
                    childrenNodes.add(treeNode);
                    List<TreeNode<T>> subNodes = allChildrenNodes(treeNodes, treeNode.getId());
                    if (subNodes != null && !subNodes.isEmpty()) {
                        childrenNodes.addAll(subNodes);
                    }
                }
            }

            return childrenNodes;
        }
        return null;
    }

    /**
     * 扁平化数结构
     * @param all 树结构结合
     * @return List
     * @param <T> all
     */
    public static <T> List<TreeNode<T>> flattenTree(List<TreeNode<T>> all) {
        List<TreeNode<T>> result = new ArrayList<>();
        if (all != null && !all.isEmpty()) {
            for (TreeNode<T> tTreeNode : all) {
                result.add(tTreeNode);
                List<TreeNode<T>> childrenNodes = flattenTree(tTreeNode.getChildren());
                result.addAll(childrenNodes);
            }
        }
        return result;
    }

    /**
     * 原始节点转换成TreeNode节点
     * @param sourceNodes 原始节点集合
     * @param idFunc id函数对象
     * @param pidFunc pid函数对象
     * @param nameFunc name函数对象
     * @return List<TreeNode>
     * @param <T> 泛型
     */
    public static <T> List<TreeNode<T>> convertToTreeNodes(List<T> sourceNodes, Function<T, Long> idFunc, Function<T, Long> pidFunc, Function<T, String> nameFunc) {
        if (sourceNodes != null && !sourceNodes.isEmpty()) {
            return sourceNodes.stream().map(t -> {
                TreeNode<T> treeNode = new TreeNode<>();
                treeNode.setTarget(t);
                treeNode.setId(idFunc.apply(t));
                treeNode.setPid(pidFunc.apply(t));
                treeNode.setName(nameFunc.apply(t));

                return treeNode;
            }).collect(Collectors.toList());
        }
        return null;
    }

    private static <T> List<TreeNode<T>> getChildren(TreeNode<T> parentTreeNode, List<TreeNode<T>> all) {
        return all.stream()
                .filter(tTreeNode -> Objects.equals(parentTreeNode.getId(), tTreeNode.getPid()))
                .peek(tTreeNode -> tTreeNode.setChildren(getChildren(tTreeNode, all))).collect(Collectors.toList());
    }
}
